﻿using Microsoft.ML.OnnxRuntime;
using Microsoft.ML.OnnxRuntime.Tensors;
using System;
using System.Collections.Generic;
using System.Drawing.Imaging;
using System.Drawing;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading.Tasks;
using System.Windows.Forms;

namespace VisionRoc.Core
{
    public class Yolo_Detection
    {
        private InferenceSession inference;
        private int tensorW, tensorH;
        private string inputName;
        private int[] inputTensorData;
        private DenseTensor<float> inputTensor;
        private float zoomRatio = 1;
        public string[] labels { get; set; }

        /// <summary>
        /// 初始化推理模型
        /// </summary>
        /// <param name="modelPath">onnx路径</param>
        /// <param name="GPUindex">启用GPU从0开始的的索引,该索引包含集显(如果有的话),需自行测试要使用的GPU索引,-1为使用CPU</param>
        public Yolo_Detection(string modelPath, int GPUindex = 0)
        {
            try  
            {
                if (GPUindex > -1)
                {
                    SessionOptions mode = new SessionOptions();
                    mode.AppendExecutionProvider_DML(GPUindex);
                    inference = new InferenceSession(modelPath, mode);
                }
                else
                {
                    inference = new InferenceSession(modelPath);
                }

                inputName = inference.InputNames.First();//images

                inputTensorData = inference.InputMetadata[inputName].Dimensions;//如{ 1,3,640,640}                                                         
                var modelData = inference.ModelMetadata.CustomMetadataMap;

                labels = SplitLabels(modelData["names"]);
                tensorW = inputTensorData[3];
                tensorH = inputTensorData[2];

            }
            catch (Exception ex)
            {
                MessageBox.Show(ex.Message, "错误", MessageBoxButtons.OK, MessageBoxIcon.Error);

            }

        }

        /// <summary>
        /// 推理主方法
        /// </summary>
        /// <param name="pic">要推理的图片</param>
        /// <param name="credibility">可信度0-1的值</param>
        /// <param name="iouThreshold">iou阈值,表示允许最大重叠比例</param>
        /// <param name="iouOverall">是否全局iou,表示不同类型是否也按iou进行计算</param>
        /// <returns></returns>
        public List<YOLOData> Detect(Bitmap pic, float credibility = 0.5f, float iouThreshold = 0.3f, bool iouOverall = false)
        {
            inputTensor = new DenseTensor<float>(inputTensorData);
            zoomRatio = 1;
            inputTensor = writeTensor(pic, inputTensorData);
            IReadOnlyCollection<NamedOnnxValue> container = new List<NamedOnnxValue> { NamedOnnxValue.CreateFromTensor(inputName, inputTensor) };
            Tensor<float> output0 = inference.Run(container).First().AsTensor<float>();
            List<YOLOData> filterData = CredibilityFiltering(output0, credibility);
            List<YOLOData> lastData = NMS(filterData, iouThreshold, iouOverall);
            Restore(ref lastData);
            return lastData;
        }

        /// <summary>
        /// 生成推理图像
        /// </summary>
        /// <param name="pic">原始图像</param>
        /// <param name="data">推理数据</param>
        /// <param name="labels">标签组,可使用模型预设的标签组,如yolo.labels</param>
        /// <returns></returns>
        public Image GenerateImage(Image pic, List<YOLOData> data, string[] labels)
        {
            Bitmap reImage = new Bitmap(pic.Width, pic.Height);
            Graphics g = Graphics.FromImage(reImage);
            g.TextRenderingHint = System.Drawing.Text.TextRenderingHint.AntiAlias;

            int penW = (pic.Width > pic.Height ? pic.Height : pic.Width) / 135;
            penW = penW <= 0 ? 1 : penW;
            Pen boxPen = new Pen(Color.BlueViolet, penW);

            int fontW = (pic.Width > pic.Height ? pic.Height : pic.Width) / 38;
            fontW = fontW <= 0 ? 1 : fontW;
            Font font = new Font("宋体", fontW, FontStyle.Bold);

            SolidBrush textColor = new SolidBrush(Color.Black);
            SolidBrush backdropColor = new SolidBrush(Color.Orange);
            float textW;
            float textH;

            g.DrawImage(pic, 0, 0, pic.Width, pic.Height);
            string writeText;

            ConvertDrawing(ref data);


            for (int i = 0; i < data.Count; i++)
            {

                string credibility = data[i].basic[4].ToString("_0.00");
                if ((int)data[i].basic[5] + 1 > labels.Length)
                {
                    writeText = credibility;
                }
                else
                {
                    writeText = labels[(int)data[i].basic[5]] + credibility;
                }
                textW = g.MeasureString(writeText + "_0.00", font).Width;
                textH = g.MeasureString(writeText + "_0.00", font).Height;
                Rectangle rectangle = new Rectangle((int)data[i].basic[0], (int)data[i].basic[1], (int)data[i].basic[2], (int)data[i].basic[3]);
                g.DrawRectangle(boxPen, rectangle);
                g.FillRectangle(backdropColor, data[i].basic[0] - boxPen.Width / 2 - 1, data[i].basic[1] - textH - boxPen.Width / 2 - 1, textW * 0.8f, textH);
                g.DrawString(writeText, font, textColor, data[i].basic[0] - boxPen.Width / 2 - 1, data[i].basic[1] - textH - boxPen.Width / 2 - 1);
            }


            ConvertCenter(ref data);
            g.Dispose();
            return reImage;

        }


        private DenseTensor<float> writeTensor(Bitmap pic, int[] inputTensorData)
        {
            BitmapData picData = pic.LockBits(new Rectangle(0, 0, pic.Width, pic.Height), ImageLockMode.ReadOnly, PixelFormat.Format24bppRgb);
            int stride = picData.Stride;
            IntPtr scan0 = picData.Scan0;
            float[,,] tempData = new float[inputTensorData[1], inputTensorData[2], inputTensorData[3]];
            float ZoomPicW = pic.Width;
            float ZoomPicH = pic.Height;

            if (ZoomPicW > tensorW || ZoomPicH > tensorH)
            {
                zoomRatio = (tensorW / ZoomPicW) < (tensorH / ZoomPicH) ? (tensorW / ZoomPicW) : (tensorH / ZoomPicH);
                ZoomPicW = ZoomPicW * zoomRatio;
                ZoomPicH = ZoomPicH * zoomRatio;
            }
            int XLocation, YLocation;
            float coefficient = 1 / zoomRatio;

            for (int y = 0; y < (int)ZoomPicH; y++)
            {
                for (int x = 0; x < (int)ZoomPicW; x++)
                {
                    XLocation = (int)(x * coefficient);
                    YLocation = (int)(y * coefficient);

                    IntPtr pixel = IntPtr.Add(scan0, YLocation * stride + XLocation * 3);
                    tempData[2, y, x] = Marshal.ReadByte(pixel) / 255f;

                    pixel = IntPtr.Add(pixel, 1);
                    tempData[1, y, x] = Marshal.ReadByte(pixel) / 255f;

                    pixel = IntPtr.Add(pixel, 1);
                    tempData[0, y, x] = Marshal.ReadByte(pixel) / 255f;

                }

            }

            pic.UnlockBits(picData);
            float[] unfoldTempData = new float[inputTensorData[1] * inputTensorData[2] * inputTensorData[3]];
            Buffer.BlockCopy(tempData, 0, unfoldTempData, 0, unfoldTempData.Length * 4);
            return new DenseTensor<float>(unfoldTempData, inputTensorData);


        }
        private List<YOLOData> CredibilityFiltering(Tensor<float> data, float credibility)
        {
            List<YOLOData> REData = new List<YOLOData>();
            int outDataW = data.Dimensions[2];
            float tempCredibility;
            int index;
            float[] data2 = data.ToArray();
            for (int i = 0; i < data2.Length; i += outDataW)
            {
                tempCredibility = 0f;
                index = -1;
                for (int j = 0; j < outDataW - 4; j++)
                {
                    if (data2[i + 4 + j] > credibility)
                    {
                        if (tempCredibility < data2[i + 4 + j])
                        {
                            tempCredibility = data2[i + 4 + j];
                            index = j;
                        }
                    }
                }
                if (index != -1)
                {
                    float[] tempData = new float[6];
                    YOLOData temp = new YOLOData();
                    tempData[0] = data2[i];
                    tempData[1] = data2[i + 1];
                    tempData[2] = data2[i + 2];
                    tempData[3] = data2[i + 3];
                    tempData[4] = tempCredibility;
                    tempData[5] = index;
                    temp.basic = tempData;
                    REData.Add(temp);
                }
            }
            return REData;
        }
        private List<YOLOData> NMS(List<YOLOData> firstData, float IOU, bool iouOverall)
        {
            CredibilitySort(ref firstData);
            List<YOLOData> nmsArray = new List<YOLOData>();
            bool satisfy = true;
            for (int i = 0; i < firstData.Count; i++)
            {
                for (int j = 0; j < nmsArray.Count; j++)
                {
                    if (iouOverall || firstData[i].basic[5] == nmsArray[j].basic[5])
                    {
                        float a = CalculatingIOU(firstData[i].basic, nmsArray[j].basic);
                        if (a > IOU)
                        {
                            satisfy = false;
                            break;
                        }
                        else
                        {
                            satisfy = true;
                        }
                    }
                    else
                    {
                        satisfy = true;
                    }
                }
                if (satisfy) nmsArray.Add(firstData[i]);
            }
            return nmsArray;
        }
        private void CredibilitySort(ref List<YOLOData> data)
        {
            int n = data.Count;
            if (n > 0)
            {
                for (int i = 0; i < n - 1; i++)
                {
                    for (int j = 0; j < n - i - 1; j++)
                    {

                        if (data[j].basic[4] < data[j + 1].basic[4])
                        {
                            YOLOData temp = data[j];
                            data[j] = data[j + 1];
                            data[j + 1] = temp;
                        }
                    }
                }
            }
        }
        private void Restore(ref List<YOLOData> data)
        {

            if (data.Count > 0)
            {

                for (int i = 0; i < data.Count; i++)
                {
                    data[i].basic[0] = data[i].basic[0] / zoomRatio;
                    data[i].basic[1] = data[i].basic[1] / zoomRatio;
                    data[i].basic[2] = data[i].basic[2] / zoomRatio;
                    data[i].basic[3] = data[i].basic[3] / zoomRatio;
                }

            }
        }
        private float CalculatingIOU(float[] rectangle1, float[] rectangle2)
        {
            float[] rectangle3 = new float[4];
            float[] rectangle4 = new float[4];

            rectangle3[0] = rectangle1[0] - rectangle1[2] / 2;
            rectangle3[1] = rectangle1[1] - rectangle1[3] / 2;
            rectangle3[2] = rectangle1[0] + rectangle1[2] / 2;
            rectangle3[3] = rectangle1[1] + rectangle1[3] / 2;

            rectangle4[0] = rectangle2[0] - rectangle2[2] / 2;
            rectangle4[1] = rectangle2[1] - rectangle2[3] / 2;
            rectangle4[2] = rectangle2[0] + rectangle2[2] / 2;
            rectangle4[3] = rectangle2[1] + rectangle2[3] / 2;

            float IntersectionArea, UnionArea;
            float left = Math.Max(rectangle3[0], rectangle4[0]);
            float up = Math.Max(rectangle3[1], rectangle4[1]);
            float right = Math.Min(rectangle3[2], rectangle4[2]);
            float down = Math.Min(rectangle3[3], rectangle4[3]);

            if (left < right && up < down)
            {
                IntersectionArea = (right - left) * (down - up);
            }
            else
            {
                IntersectionArea = 0;
            }

            float area1 = (rectangle3[2] - rectangle3[0]) * (rectangle3[3] - rectangle3[1]);
            float area2 = (rectangle4[2] - rectangle4[0]) * (rectangle4[3] - rectangle4[1]);
            UnionArea = area1 + area2 - IntersectionArea;
            return IntersectionArea / UnionArea;
        }
        private string[] SplitLabels(string labels)
        {
            string temp1 = labels.Replace("{", "").Replace("}", "");
            string[] tempArray = temp1.Split(',');
            string[] reArray = new string[tempArray.Length];
            for (int i = 0; i < tempArray.Length; i++)
            {
                int SpStart = tempArray[i].IndexOf(':') + 3;
                int SpEnd = tempArray[i].Length - 1;
                reArray[i] = tempArray[i].Substring(SpStart, SpEnd - SpStart);
            }
            return reArray;
        }
     
        private void ConvertDrawing(ref List<YOLOData> data)
        {

            if (data.Count > 0 && data[0].basic.Length > 2)
            {
                for (int i = 0; i < data.Count; i++)
                {
                    data[i].basic[0] = data[i].basic[0] - data[i].basic[2] / 2;
                    data[i].basic[1] = data[i].basic[1] - data[i].basic[3] / 2;
                }
            }


        }
        private void ConvertCenter(ref List<YOLOData> data)
        {
            //还原中心坐标
            if (data.Count > 0 && data[0].basic.Length > 2)
            {
                for (int i = 0; i < data.Count; i++)
                {
                    data[i].basic[0] = data[i].basic[0] + data[i].basic[2] / 2;
                    data[i].basic[1] = data[i].basic[1] + data[i].basic[3] / 2;
                }
            }

        }


    }
}
